import cv2
import numpy as np
from utils.common import load_image, make_dirs  # 导入公共函数
from utils.config import exp4_drainage_path, exp4_output_dir  # 导入输出路径

def watershed_segmentation(image):
    """
    使用分水岭算法进行图像分割
    :param image: 输入彩色图像
    :return: 分割后的图像
    """
    # 将图像转换为灰度图像
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 二值化处理
    ret, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 去除噪声
    kernel = np.ones((3, 3), np.uint8)
    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)

    # 确定背景区域
    sure_bg = cv2.dilate(opening, kernel, iterations=3)

    # 寻找前景区域
    dist_transform = cv2.distanceTransform(opening, cv2.DIST_L2, 5)
    ret, sure_fg = cv2.threshold(dist_transform, 0.7 * dist_transform.max(), 255, 0)

    # 确定未知区域
    sure_fg = np.uint8(sure_fg)
    unknown = cv2.subtract(sure_bg, sure_fg)

    # 标记标签
    ret, markers = cv2.connectedComponents(sure_fg)
    markers += 1
    markers[unknown == 255] = 0

    # 应用分水岭算法
    markers = cv2.watershed(image, markers)

    # 标记分割边界
    image[markers == -1] = [255, 0, 0]  # 将边界标记为红色

    return image

def main():
    # 确保输出目录存在
    make_dirs(exp4_output_dir)

    # 加载图片
    image = load_image()  # 使用公共函数加载图像

    # 使用分水岭算法进行图像分割
    segmented_image = watershed_segmentation(image)

    # 显示结果
    cv2.imshow("Original Image", image)
    cv2.imshow("Watershed Segmentation", segmented_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 保存结果
    cv2.imwrite(f"{exp4_output_dir}/watershed_result.jpg", segmented_image)

if __name__ == "__main__":
    main()